# -*- coding: utf-8 -*-

"""
工具函数
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import unicode_literals

import io
import json
import os
import sys
import hashlib
from copy import deepcopy
from importlib import import_module
from zipimport import zipimporter, ZipImportError

import six
from six.moves import cPickle
from typing import Optional, AnyStr, Any, Dict, List, Tuple

from kgpipeline.jsonutil import default_encode

import logging

logger = logging.getLogger(__name__)


def compare_3_way(a, b):
    """
    We don't have cmp function in Python3, so use this as a compatibility layer
    """
    if a > b:
        return 1
    elif a < b:
        return -1
    else:
        return 0


class IdUtil(object):

    HASH_FUNCTION = hashlib.sha256

    @classmethod
    def any2hash(cls, text):
        """
        返回string或者其他对象的sha256 hash值

        :param text: str, list, dict。
            如果传入参数是list、dict，则按下列方式预处理为字符串:
            - list, 使用 ``___``连接
            - dict, 转换为json字符串
        :return: hash : str。sha256 hash, 16进制字符串
        """
        # optimize list of text to speedup
        try:
            assert text
            if isinstance(text, list):
                if len(text) == 1:
                    text = text[0]
                else:
                    text = "___".join(text)
        except TypeError:
            pass

        if isinstance(text, (dict, list)):
            text = json.dumps(text, sort_keys=True, default=default_encode)

        # if unicode then encode as utf8
        if isinstance(text, six.text_type):
            text = text.encode('utf-8')

        # if text is still not bytes, then convert to bytes forcibly
        if not isinstance(text, six.binary_type):
            text = six.binary_type(text)

        return cls.HASH_FUNCTION(text).hexdigest()

    @classmethod
    def create_dummy_entity(cls, uid, _type=None):
        """
        创建空entity，用于link

        :param uid: uid
        :param _type:
        :return: entity
        """
        if _type is None:
            return {"@id": cls.any2hash(uid)}
        elif isinstance(uid, list):
            uid = deepcopy(uid)
            uid.append(_type)
            return {"@id": cls.any2hash(uid)}
        else:
            return {"@id": cls.any2hash([six.text_type(uid), _type])}


def import_direct_from_module(path):  # type: (AnyStr) -> Any
    """
    :param path: str，module的路径
    :return: type, imported module
    :raises: ImportError，import失败时
    """
    """
    import module

    Parameters
    ----------
    path : str
        module path

    Returns
    -------
    module : type
        imported module

    Raises
    ------
    ImportError:
        when fail to import
    """
    path_list = path.rsplit(".", 1)
    if len(path_list) == 2:
        class_name = path_list[1]
        module_path = path_list[0]
        m = import_module(module_path)
        if hasattr(m, class_name):
            class_type = getattr(m, class_name)
            return class_type
        else:
            raise ImportError("No member {0} defined in module {1}".format(class_name, module_path))
    raise ImportError("Malformed module path: {0}".format(path))


def import_class(path, super_type=None):  # type: (AnyStr, Optional[type]) -> type
    """
    import class (and check its super class)

    Parameters
    ----------
    path : str
        module path
    super_type : type, optional
        if provided, will need that it is the super class of imported class.

    Returns
    -------
    class : type
        the imported class

    Raises
    ------
    ImportError
        when not imported or imported class is not subclass of super_type
    """
    class_type = import_direct_from_module(path)
    if not isinstance(class_type, type):
        raise ImportError("{0} is not a class".format(class_type))
    if super_type is not None and super_type not in class_type.mro():
        raise ImportError("{0} is not a subclass of {1}".format(class_type.__name__, super_type.__name__))
    return class_type


def wrap_broadcast_variable(mapper, schema):
    def __func(row):
        return mapper(schema.value, row)

    return __func


def load_pickle(path):
    with io.open(path, "rb") as f:
        return cPickle.load(f)


def save_pickle(obj, path):
    with io.open(path, "wb") as f:
        cPickle.dump(obj, f, cPickle.HIGHEST_PROTOCOL)


def load_json(path):
    with io.open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_json(d, path, **kwargs):
    """
    convert dict into json and save it to file in utf-8.

    Parameters
    ----------
    d : dict
        the value of json in dict format
    path : str
        json file path
    kwargs:
        arguments passed to json.dump()

    Notes
    -----
    if parameter `ensure_ascii` is not specified, will append this parameter with value `False`

    Returns
    -------
    None
    """
    if 'ensure_ascii' not in kwargs:
        kwargs['ensure_ascii'] = False

    with io.open(path, "w+", encoding="utf-8") as f:
        c = json.dumps(d, **kwargs)
        f.write(c)


# region Load data from module path

def load_schema_from_module_path(path):
    b = load_data_from_module_path(path)
    if path.endswith('.json'):
        return json.loads(b.decode('utf-8'))
    else:
        return cPickle.loads(b)


def load_schema(path):
    if path.endswith(".json"):
        return json.loads(load_text_from_module_or_fs_path(path))
    elif path.endswith(".pkl"):
        return cPickle.loads(load_data_from_module_or_fs_path(path))
    else:
        raise Exception("Unsupported schema file format: {0}!".format(path))


def load_data_from_module_path(path):
    """
    read binary data from file in module (or zipped) in the search path

    Parameters
    ----------
    path : str
        file path

    Returns
    -------
    bytes:
        the binary content in file

    Raises
    ------
    IOError
        when fail to open
    """
    for p in sys.path:
        if p.endswith(".zip"):
            try:
                zi = zipimporter(p)
                return zi.get_data(path)
            except ZipImportError:
                continue
            except IOError:
                logger.warning("can not access {} when importing {}".format(p, path))

        else:
            full_path = os.path.join(p, path)
            if os.path.exists(full_path):
                with open(full_path, 'rb') as f:
                    return f.read()
    raise IOError


def load_data_from_module_or_fs_path(path):
    """
    read binary data from file in module (or zipped) in the search path or local file system path

    Parameters
    ----------
    path : str
        file path

    Returns
    -------
    data:
        byte string
    """
    try:
        data = load_data_from_module_path(path)
    except IOError:
        if os.path.exists(path):
            with io.open(path, "rb") as f:
                data = f.read()
        else:
            raise Exception("Path {0} is not a valid module path or file system path!".format(path))
    return data


def load_text_from_module_path(path):
    """
    read text data from file in module (or zipped) in the search path

    Parameters
    ----------
    path : str
        file path

    Returns
    -------
    str:
        the text content in utf-8 in file
    """
    data = load_data_from_module_path(path)
    s = data.decode('utf-8')
    return s


def load_text_from_module_or_fs_path(path):
    """
    read text data from file in module (or zipped) in the search path or local file system path

    Parameters
    ----------
    path : str
        file path

    Returns
    -------
    str:
        the text content in utf-8 in file
    """
    try:
        data = load_data_from_module_path(path)
        s = data.decode('utf-8')
    except IOError:
        if os.path.exists(path):
            with io.open(path, "r", encoding="utf-8") as f:
                s = f.read()
        else:
            raise Exception("Path {0} is not a valid module path or file system path!".format(path))
    return s


# endregion


# region data structure utilities

def generate_dag_sequence(dependency_dict):  # type: (Dict[AnyStr, List[AnyStr]]) -> List[Tuple[AnyStr,bool]]
    """
    produce a linear sequence for a DAG graph
    :param dependency_dict: {id => [dep_id, dep_id, ...]}
    :return: [(id, if being depended), ...]
    """
    dep = {k: set(v) for k, v in dependency_dict.items()}
    inv_dep = {}
    for k, v in dep.items():
        for sid in v:
            if sid not in dep:
                raise Exception("ID {0} not in total set!".format(sid))
            if sid not in inv_dep:
                inv_dep[sid] = {k}
            else:
                inv_dep[sid].add(k)

    processed = set()
    res = []

    while True:
        last_processed = len(processed)
        for k, v in dep.items():
            if not v and k not in processed:
                processed.add(k)
                if k in inv_dep:
                    for sid in inv_dep[k]:
                        dep[sid].remove(k)
                    del inv_dep[k]
                    res.append((k, True))
                else:
                    res.append((k, False))
        if len(processed) == len(dep):
            break
        elif last_processed == len(processed):
            raise Exception("No progress! There is potential infinite loop!")

    return res


# endregion


def escape_str(s):
    """
    convert text readable escaped string into working escaped, like '\\t' to '\t' and '\\n' to '\n'.

    Parameters
    ----------
    s : str
        input string with readable escaped string

    Returns
    -------
    str:
        escaped string

    Examples
    --------
    '\\t' -> '\t'
    '\x09' -> '\t'
    '\\n' -> '\n'

    Notes
    -----
    in python 2.7, it can be implemented as:
        s.encode('utf-8').decode('string-escape').decode('utf-8')

    in python 3.7, following code will fail generally on utf-8 (eg: '你好\\n'):
        s.encode('utf-8').decode('unicode-escape')

    References
    ----------
        1. https://stackoverflow.com/questions/4020539/process-escape-sequences-in-a-string-in-python/24519338
    """
    s1 = s.encode('raw-unicode-escape').decode('unicode-escape')
    return s1


def by_batch(items, batch_size):
    """
    return items in batches
    """
    batch_count = 0
    batch = []
    for item in items:
        batch_count += 1
        batch.append(item)
        if batch_count > batch_size:
            yield batch
            batch_count = 0
            batch = []
    if batch_count > 0:
        yield batch
